from flask import Flask, Response
import cv2
import torch

app = Flask(__name__)

# 加载你训练好的YOLOv5模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', source='local')
model.conf = 0.25  # 设置置信度阈值（可选）

# 打开摄像头（0）或视频文件路径
cap = cv2.VideoCapture(0)

def gen_frames():
    while True:
        success, frame = cap.read()
        if not success:
            break

        # YOLOv5 推理
        results = model(frame)

        # 渲染结果（会在原图上绘制检测框）
        img = results.render()[0]

        # 编码为JPEG格式
        ret, buffer = cv2.imencode('.jpg', img)
        if not ret:
            continue
        frame_bytes = buffer.tobytes()

        # 推送为MJPEG流
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

@app.route('/video_feed')
def video_feed():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
